"""
直方图计算工具模块

该模块提供颜色直方图计算和分析功能。
主要功能：
1. 颜色直方图计算
2. 直方图比较和相似度计算
3. 直方图可视化
4. 颜色空间转换

作者：目标跟踪系统
日期：2024
"""

import cv2
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Optional, List


class HistogramCalculator:
    """
    直方图计算器类
    
    提供各种颜色空间的直方图计算和分析功能。
    支持HSV、RGB、LAB等颜色空间。
    
    参数：
        color_space: 颜色空间，默认为'HSV'
        bins: 直方图分箱数，默认为[16, 16]
        ranges: 颜色值范围，默认为None（自动设置）
    """
    
    def __init__(self, color_space: str = 'HSV', bins: List[int] = [16, 16], ranges: Optional[List[int]] = None):
        self.color_space = color_space.upper()
        self.bins = bins
        self.ranges = ranges or self._get_default_ranges()
        
    def _get_default_ranges(self) -> List[int]:
        """
        获取默认的颜色值范围
        
        返回：
            颜色值范围列表
        """
        if self.color_space == 'HSV':
            return [0, 180, 0, 256]  # H: 0-180, S: 0-256
        elif self.color_space == 'RGB':
            return [0, 256, 0, 256, 0, 256]  # R: 0-256, G: 0-256, B: 0-256
        elif self.color_space == 'LAB':
            return [0, 256, 0, 256, 0, 256]  # L: 0-256, A: 0-256, B: 0-256
        else:
            return [0, 256, 0, 256]  # 默认范围
            
    def convert_color_space(self, image: np.ndarray) -> np.ndarray:
        """
        转换颜色空间
        
        参数：
            image: 输入图像（BGR格式）
            
        返回：
            转换后的图像
        """
        if self.color_space == 'HSV':
            return cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        elif self.color_space == 'RGB':
            return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        elif self.color_space == 'LAB':
            return cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        else:
            return image
            
    def calculate_histogram(self, image: np.ndarray, mask: Optional[np.ndarray] = None) -> np.ndarray:
        """
        计算图像的颜色直方图
        
        参数：
            image: 输入图像
            mask: 掩码图像，如果提供则只计算掩码区域的直方图
            
        返回：
            归一化的直方图
        """
        # 转换颜色空间
        converted_image = self.convert_color_space(image)
        
        # 选择通道
        if self.color_space in ['HSV', 'RGB', 'LAB']:
            channels = [0, 1]  # 使用前两个通道
        else:
            channels = [0, 1] if converted_image.ndim == 3 else [0]
            
        # 计算直方图
        hist = cv2.calcHist([converted_image], channels, mask, self.bins, self.ranges)
        
        # 归一化
        cv2.normalize(hist, hist, 0, 255, cv2.NORM_MINMAX)
        
        return hist
        
    def calculate_roi_histogram(self, image: np.ndarray, roi: Tuple[int, int, int, int]) -> np.ndarray:
        """
        计算感兴趣区域(ROI)的直方图
        
        参数：
            image: 输入图像
            roi: 感兴趣区域 (x, y, width, height)
            
        返回：
            ROI区域的直方图
        """
        x, y, w, h = roi
        
        # 检查ROI是否在图像范围内
        if x < 0 or y < 0 or x + w > image.shape[1] or y + h > image.shape[0]:
            raise ValueError("ROI超出图像范围")
            
        # 提取ROI
        roi_image = image[y:y+h, x:x+w]
        
        return self.calculate_histogram(roi_image)
        
    def calculate_masked_histogram(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """
        计算带掩码的直方图
        
        参数：
            image: 输入图像
            mask: 掩码图像
            
        返回：
            掩码区域的直方图
        """
        return self.calculate_histogram(image, mask)
        
    def compare_histograms(self, hist1: np.ndarray, hist2: np.ndarray, method: str = 'correlation') -> float:
        """
        比较两个直方图的相似度
        
        参数：
            hist1: 第一个直方图
            hist2: 第二个直方图
            method: 比较方法，可选：'correlation', 'chi_square', 'intersection', 'bhattacharyya'
            
        返回：
            相似度值
        """
        method_map = {
            'correlation': cv2.HISTCMP_CORREL,
            'chi_square': cv2.HISTCMP_CHISQR,
            'intersection': cv2.HISTCMP_INTERSECT,
            'bhattacharyya': cv2.HISTCMP_BHATTACHARYYA
        }
        
        if method not in method_map:
            raise ValueError(f"不支持的比较方法: {method}")
            
        return cv2.compareHist(hist1, hist2, method_map[method])
        
    def calculate_similarity(self, hist1: np.ndarray, hist2: np.ndarray) -> float:
        """
        计算两个直方图的相似度（0-1范围）
        
        参数：
            hist1: 第一个直方图
            hist2: 第二个直方图
            
        返回：
            相似度值，1表示完全相似，0表示完全不同
        """
        # 使用Bhattacharyya距离计算相似度
        distance = self.compare_histograms(hist1, hist2, 'bhattacharyya')
        
        # 转换为相似度（距离越小，相似度越高）
        similarity = 1.0 - distance
        
        return max(0.0, min(1.0, similarity))
        
    def create_backprojection(self, image: np.ndarray, hist: np.ndarray) -> np.ndarray:
        """
        创建反向投影图像
        
        参数：
            image: 输入图像
            hist: 目标直方图
            
        返回：
            反向投影图像
        """
        # 转换颜色空间
        converted_image = self.convert_color_space(image)
        
        # 选择通道
        if self.color_space in ['HSV', 'RGB', 'LAB']:
            channels = [0, 1]
        else:
            channels = [0, 1] if converted_image.ndim == 3 else [0]
            
        # 计算反向投影
        backproj = cv2.calcBackProject([converted_image], channels, hist, self.ranges, 1)
        
        return backproj
        
    def visualize_histogram(self, hist: np.ndarray, title: str = "直方图") -> None:
        """
        可视化直方图
        
        参数：
            hist: 要可视化的直方图
            title: 图像标题
        """
        plt.figure(figsize=(10, 6))
        
        if len(self.bins) == 1:
            # 一维直方图
            plt.plot(hist)
            plt.xlabel('颜色值')
            plt.ylabel('频率')
        else:
            # 二维直方图
            plt.imshow(hist, interpolation='nearest', cmap='hot')
            plt.colorbar()
            plt.xlabel('颜色通道1')
            plt.ylabel('颜色通道2')
            
        plt.title(title)
        plt.show()
        
    def analyze_histogram(self, hist: np.ndarray) -> dict:
        """
        分析直方图统计信息
        
        参数：
            hist: 输入直方图
            
        返回：
            包含统计信息的字典
        """
        # 计算统计信息
        total_pixels = np.sum(hist)
        mean_value = np.mean(hist)
        std_value = np.std(hist)
        max_value = np.max(hist)
        min_value = np.min(hist)
        
        # 计算熵（信息量）
        normalized_hist = hist / total_pixels
        entropy = -np.sum(normalized_hist * np.log2(normalized_hist + 1e-10))
        
        # 计算主要颜色分布
        flat_hist = hist.flatten()
        sorted_indices = np.argsort(flat_hist)[::-1]
        dominant_bins = sorted_indices[:5]  # 前5个主要颜色
        
        return {
            'total_pixels': total_pixels,
            'mean': mean_value,
            'std': std_value,
            'max': max_value,
            'min': min_value,
            'entropy': entropy,
            'dominant_bins': dominant_bins.tolist(),
            'dominant_values': flat_hist[dominant_bins].tolist()
        }
        
    def adaptive_histogram_equalization(self, image: np.ndarray, clip_limit: float = 2.0, 
                                      tile_grid_size: Tuple[int, int] = (8, 8)) -> np.ndarray:
        """
        自适应直方图均衡化
        
        参数：
            image: 输入图像
            clip_limit: 对比度限制阈值
            tile_grid_size: 瓦片网格大小
            
        返回：
            均衡化后的图像
        """
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
            
        # 创建CLAHE对象
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
        
        # 应用自适应直方图均衡化
        equalized = clahe.apply(gray)
        
        return equalized
        
    def match_histograms(self, source: np.ndarray, reference: np.ndarray) -> np.ndarray:
        """
        直方图匹配
        
        将源图像的直方图匹配到参考图像的直方图。
        
        参数：
            source: 源图像
            reference: 参考图像
            
        返回：
            匹配后的图像
        """
        # 计算累积分布函数
        def calculate_cdf(hist):
            cdf = hist.cumsum()
            cdf_normalized = cdf / cdf[-1]
            return cdf_normalized
            
        # 转换为灰度图像
        if len(source.shape) == 3:
            source_gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY)
        else:
            source_gray = source
            
        if len(reference.shape) == 3:
            reference_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY)
        else:
            reference_gray = reference
            
        # 计算直方图
        source_hist = cv2.calcHist([source_gray], [0], None, [256], [0, 256])
        reference_hist = cv2.calcHist([reference_gray], [0], None, [256], [0, 256])
        
        # 计算累积分布函数
        source_cdf = calculate_cdf(source_hist)
        reference_cdf = calculate_cdf(reference_hist)
        
        # 创建映射表
        mapping = np.zeros(256, dtype=np.uint8)
        for i in range(256):
            # 找到最接近的参考值
            diff = np.abs(reference_cdf - source_cdf[i])
            mapping[i] = np.argmin(diff)
            
        # 应用映射
        matched = cv2.LUT(source_gray, mapping)
        
        return matched 